"""硬件扫描模块，用于收集系统信息"""

import platform
import re

import psutil
from ocl_llm.utils.display import display

# 尝试导入pynvml库（在模块顶部导入）
try:
    from pynvml import (
        nvmlInit,
        nvmlDeviceGetCount,
        nvmlDeviceGetHandleByIndex,
        nvmlDeviceGetMemoryInfo,
        nvmlDeviceGetName,
        NVMLError,
        NVML_ERROR_LIBRARY_NOT_FOUND,
    )

    HAS_NVML = True
except ImportError:
    HAS_NVML = False
    # 创建虚拟的NVMLError占位符和错误码
    NVML_ERROR_LIBRARY_NOT_FOUND = 15

    class NVMLError(Exception):
        """NVML错误占位符"""


class HardwareScanner:
    """扫描并报告系统硬件信息"""

    def __init__(self):
        """
        功能：初始化硬件扫描器
        :return: 无
        """
        self.cpu_info = {}
        self.gpu_info = {}
        self.memory = 0
        self.disk_info = {}

    def scan(self):
        """
        功能：执行完整的硬件扫描
        :return: 包含硬件信息的字典
        """
        self.scan_cpu()
        self.scan_memory()
        self.scan_gpu()
        self.scan_disk()
        return {
            "cpu": self.cpu_info,
            "memory": self.memory,
            "gpu": self.gpu_info,
            "disk": self.disk_info,
        }

    def scan_cpu(self):
        """
        功能：扫描CPU信息并计算得分
        :return: 无
        """
        try:
            # 获取CPU核心和线程数
            cores = psutil.cpu_count(logical=False) or 1
            threads = psutil.cpu_count(logical=True) or 2

            # 获取CPU基础频率
            base_freq = 3.0
            try:
                freq_info = psutil.cpu_freq()
                if freq_info and freq_info.max > 0:
                    base_freq = freq_info.max / 1000
                elif platform.system() == "Linux":
                    with open("/proc/cpuinfo", encoding="utf-8") as f:
                        cpuinfo = f.read()
                    freq_match = re.search(r"cpu MHz\s*:\s*(\d+\.\d+)", cpuinfo)
                    if freq_match:
                        base_freq = float(freq_match.group(1)) / 1000
            except (OSError, ValueError, TypeError):
                pass
            finally:
                base_freq = round(base_freq, 2)

            # 计算CPU得分
            base_score = cores * base_freq
            multithread_boost = 1 + (threads / cores - 1) * 0.2 if cores > 0 else 1
            score = base_score * multithread_boost

            # 获取CPU品牌信息
            brand = "未知CPU"
            if platform.system() == "Linux":
                try:
                    with open("/proc/cpuinfo", encoding="utf-8") as f:
                        cpuinfo = f.read()
                    brand_match = re.search(
                        r"model name\s*:\s*(.+)$", cpuinfo, re.MULTILINE
                    )
                    if brand_match:
                        brand = brand_match.group(1).strip()
                except (OSError, UnicodeDecodeError):
                    pass

            self.cpu_info = {
                "brand": brand,
                "cores": cores,
                "threads": threads,
                "base_freq": base_freq,
                "score": round(score, 1),
            }
        except (psutil.Error, OSError, ValueError) as e:
            display(f"⚠️ CPU扫描失败: {e}")
            self.cpu_info = {
                "brand": "未知CPU",
                "cores": 1,
                "threads": 2,
                "base_freq": 2.0,
                "score": 2.0,
            }

    def scan_memory(self):
        """
        功能：扫描系统内存
        :return: 无
        """
        try:
            self.memory = round(psutil.virtual_memory().total / (1024**3))
        except psutil.Error:
            self.memory = 8

    def scan_gpu(self):
        """
        功能：扫描GPU信息
        :return: 无
        """
        if self.try_scan_nvidia_gpu():
            return
        self.gpu_info = None

    def try_scan_nvidia_gpu(self):
        """
        功能：尝试扫描NVIDIA GPU
        :return: 如果成功扫描到GPU返回True，否则返回False
        """
        if not HAS_NVML:
            display("⚠️ 警告: pynvml模块未安装，无法检测NVIDIA GPU信息")
            return False

        try:
            # 尝试初始化NVML库
            nvmlInit()
            device_count = nvmlDeviceGetCount()
            if device_count > 0:
                handle = nvmlDeviceGetHandleByIndex(0)
                memory_info = nvmlDeviceGetMemoryInfo(handle)
                device_name = nvmlDeviceGetName(handle)

                # 处理设备名称
                if isinstance(device_name, bytes):
                    device_name = device_name.decode("utf-8")
                device_name = re.sub(r"GeForce\s*|RTX\s*|NVIDIA\s*", "", device_name)

                self.gpu_info = {
                    "name": device_name.strip(),
                    "vram": round(memory_info.total / (1024**3)),
                    "cuda_cores": self.get_cuda_cores(device_name.strip()),
                }
                return True
            return False  # 没有找到GPU设备
        except NVMLError as e:
            # 处理NVML特定错误
            if e.value == NVML_ERROR_LIBRARY_NOT_FOUND:
                display("⚠️ 警告: NVIDIA驱动未正确安装，无法扫描GPU信息")
            else:
                display(f"⚠️ 警告: NVML操作失败: {e}")
            return False
        except (OSError, RuntimeError) as e:
            display(f"⚠️ 警告: 获取NVIDIA GPU信息失败: {e}")
            return False

    def scan_disk(self):
        """
        功能：扫描磁盘空间信息
        :return: 无
        """
        try:
            disk_info = {"total": 0, "free": 0, "partitions": []}
            for partition in psutil.disk_partitions(all=False):
                try:
                    usage = psutil.disk_usage(partition.mountpoint)
                    disk_info["total"] += usage.total
                    disk_info["free"] += usage.free
                    disk_info["partitions"].append(
                        {
                            "mountpoint": partition.mountpoint,
                            "device": partition.device,
                            "total": usage.total,
                            "free": usage.free,
                        }
                    )
                except (OSError, psutil.Error):
                    continue

            disk_info["total_gb"] = round(disk_info["total"] / (1024**3))
            disk_info["free_gb"] = round(disk_info["free"] / (1024**3))
            self.disk_info = disk_info
        except (OSError, psutil.Error) as e:
            display(f"⚠️ 警告: 磁盘扫描失败: {e}")
            self.disk_info = {"total_gb": 0, "free_gb": 0, "partitions": []}

    def get_cuda_cores(self, gpu_name):
        """
        功能：基于GPU名称估算CUDA核心数
        :param gpu_name: GPU名称
        :return: 估算的CUDA核心数
        """
        clean_name = re.sub(r"[^a-zA-Z0-9\s]", "", gpu_name).strip().upper()
        clean_name = re.sub(r"\s+", " ", clean_name)

        # CUDA核心映射表
        cuda_core_map = {
            "5090 D": 16384,
            "5090": 16384,
            "4090": 16384,
            "4080": 9728,
            "3090": 10496,
            "3080": 8704,
            "3070": 5888,
            "3060": 3584,
            "3050": 2560,
            "A100": 6912,
            "A6000": 10752,
            "V100": 5120,
            "TITAN V": 5120,
            "TITAN RTX": 4608,
            "1080 TI": 3584,
            "1080": 2560,
            "1070": 1920,
            "1060": 1280,
            "1650": 896,
            "MX150": 384,
            "A5000": 8192,
            "A4000": 6144,
            "A2000": 3328,
            "QUADRO RTX 6000": 4608,
            "QUADRO RTX 8000": 4608,
            "TESLA P100": 3584,
            "TESLA T4": 2560,
        }

        # 精确匹配
        if clean_name in cuda_core_map:
            return cuda_core_map[clean_name]

        # 模式匹配
        for pattern, cores in cuda_core_map.items():
            if re.search(rf"\b{pattern}\b", clean_name):
                return cores

        # 基于型号前缀匹配
        model_match = re.search(r"(\d{4})", clean_name)
        if model_match:
            model_num = model_match.group(1)
            if model_num.startswith("50") or model_num.startswith("40"):
                return 16384
            if model_num.startswith("30"):
                return 7000

        # 基于系列匹配
        result = 0
        if "RTX 40" in clean_name or "40" in clean_name:
            result = 16384
        elif "RTX 30" in clean_name or "30" in clean_name:
            result = 10000
        elif "RTX 20" in clean_name or "20" in clean_name:
            result = 5000

        return result
